from unsloth import FastLanguageModel
import torch
max_seq_length = 2048
dtype = None
load_in_4bit = True


import wandb

wandb.login(key="822b638c1f465585a2a139ce45f36c4d0a6e0180")
run = wandb.init(
    project='my fint-tune on deepseek r1 with medical data',
    job_type="training",
    anonymous="allow"
)

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "/root/code/deepseek-gaokao/DeepSeek-R1-Distill-Llama-8B", # 这里改成你本地模型，以我的为例，我已经huggingface上的模型文件下载到本地。
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context.
Write a response that appropriately completes the request.
Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.

### Instruction:
You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning.
Please answer the following medical question.

### Question:
{}

### Response:
<think>{}"""

question = "A 61-year-old woman with a long history of involuntary urine loss during activities like coughing or sneezing but no leakage at night undergoes a gynecological exam and Q-tip test. Based on these findings, what would cystometry most likely reveal about her residual volume and detrusor contractions?"


FastLanguageModel.for_inference(model)  # Unsloth has 2x faster inference!
inputs = tokenizer([prompt_style.format(question, "")], return_tensors="pt").to("cuda")

outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=1200,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs)
print("output answer before training\r\n")
print(response[0])
print("\r\nshow response after sharpsharpsharp:\r\n")
#print(response[0].split("### Response:")[1])
# 使用 split 方法分割字符串，并提取第一个 '### Response:' 后面的所有内容
parts = response[0].split("### Response:")
if len(parts) > 1:
    # 打印第一个 '### Response:' 后面的所有内容（包括 '<think>' 和第二个 '### Response:'）
    print(parts[1:])
else:
    print("未找到 '### Response:' 标记")

FastLanguageModel.for_training(model)
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_alpha=16,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth",  # True or "unsloth" for very long context
    random_state=3407,
    use_rslora=False,
    loftq_config=None,
)
train_prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context.
Write a response that appropriately completes the request.
Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.

### Instruction:
You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning.
Please answer the following medical question.

### Question:
{}

### Response:
<think>
{}
</think>
{}"""

EOS_TOKEN = tokenizer.eos_token  # Must add EOS_TOKEN


def formatting_prompts_func(examples):
    inputs = examples["Question"]
    cots = examples["Complex_CoT"]
    outputs = examples["Response"]
    texts = []
    for input, cot, output in zip(inputs, cots, outputs):
        text = train_prompt_style.format(input, cot, output) + EOS_TOKEN
        texts.append(text)
    return {
        "text": texts,
    }

from datasets import load_dataset

dataset = load_dataset("/root/code/deepseek-gaokao/medical-o1-reasoning-SFT", "en",split = "train[0:500]") # 这里同样去huggingface上面下载数据集，然后放到本地
dataset = dataset.map(formatting_prompts_func, batched = True,)
print("\r\nfirst dataset: \r\n")
dataset["text"][0]

from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        # Use num_train_epochs = 1, warmup_ratio for full training runs!
        warmup_steps=5,
        max_steps=60,
        learning_rate=2e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=10,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
    ),
)

trainer_stats = trainer.train()

# Save the fine-tuned model
wandb.finish()

question = "A 61-year-old woman with a long history of involuntary urine loss during activities like coughing or sneezing but no leakage at night undergoes a gynecological exam and Q-tip test. Based on these findings, what would cystometry most likely reveal about her residual volume and detrusor contractions?"


FastLanguageModel.for_inference(model)  # Unsloth has 2x faster inference!
inputs = tokenizer([prompt_style.format(question, "")], return_tensors="pt").to("cuda")

outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=1200,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs)
#print("output answer after training\r\n")
#print(response[0])
print("\r\nshow response after sharpsharpsharp:\r\n")
#print(response[0].split("### Response:")[1])
# 使用 split 方法分割字符串，并提取第一个 '### Response:' 后面的所有内容
parts = response[0].split("### Response:")
if len(parts) > 1:
    # 打印第一个 '### Response:' 后面的所有内容（包括 '<think>' 和第二个 '### Response:'）
    print(parts[1:])
else:
    print("未找到 '### Response:' 标记")

question = "A 59-year-old man presents with a fever, chills, night sweats, and generalized fatigue, and is found to have a 12 mm vegetation on the aortic valve. Blood cultures indicate gram-positive, catalase-negative, gamma-hemolytic cocci in chains that do not grow in a 6.5% NaCl medium. What is the most likely predisposing factor for this patient's condition?"

inputs = tokenizer([prompt_style.format(question, "")], return_tensors="pt").to("cuda")

outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=1200,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs)
#print("output answer after training\r\n")
#print(response[0])
print("\r\nshow response after sharpsharpsharp:\r\n")
#print(response[0].split("### Response:")[1])
# 使用 split 方法分割字符串，并提取第一个 '### Response:' 后面的所有内容
parts = response[0].split("### Response:")
if len(parts) > 1:
    # 打印第一个 '### Response:' 后面的所有内容（包括 '<think>' 和第二个 '### Response:'）
    print(parts[1:])
else:
    print("未找到 '### Response:' 标记")

# new_model_online = "kingabzpro/DeepSeek-R1-Medical-COT"
new_model_local = "DeepSeek-R1-Medical-COT"
model.save_pretrained(new_model_local) # Local saving
tokenizer.save_pretrained(new_model_local)
print(f"模型和分词器已成功保存到 {new_model_local}")
model.save_pretrained_merged(new_model_local, tokenizer, save_method = "merged_16bit",)
print("all is ok")

